package db

import (
	"fmt"
	"hash/fnv"
	"log"
	"sync"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

// Sharding 动态分片管理器
type Sharding struct {
	DBs []*gorm.DB
	mu  sync.RWMutex
}

// NewSharding 创建一个新的分片管理器
func NewSharding(dbs []*gorm.DB) *Sharding {
	return &Sharding{
		DBs: dbs,
	}
}

// GetDB 根据分片键获取对应的数据库实例
func (s *Sharding) GetDB(shardKey string) *gorm.DB {
	s.mu.RLock()
	defer s.mu.RUnlock()

	hash := fnv.New32a()
	hash.Write([]byte(shardKey))
	index := hash.Sum32() % uint32(len(s.DBs))
	return s.DBs[index]
}

// InitSharding 初始化分片数据库
func InitSharding(dsnList []string) (*Sharding, error) {
	var dbs []*gorm.DB
	for _, dsn := range dsnList {
		db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
		if err != nil {
			return nil, fmt.Errorf("分片数据库连接失败: %v", err)
		}
		dbs = append(dbs, db)
	}
	log.Println("分片数据库初始化成功")
	return NewSharding(dbs), nil
}
